import os
import hashlib
import shutil
from fastapi import FastAPI, UploadFile, File
from pydantic import BaseModel
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.document_loaders import TextLoader
from langchain.docstore import InMemoryDocstore
import faiss
import requests
from vllm_request import query_vllm
import yaml
from langchain.docstore.document import Document
from fastapi.middleware.cors import CORSMiddleware
import pickle


app = FastAPI(title="RAG + vLLM Service")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 或指定 ["http://localhost:5173"]
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
# 读取配置
def load_config(config_file="config.yaml"):
    with open(config_file, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)

config = load_config()
FAISS_PATH = config["faiss"]["path"]
EMBEDDINGS_MODEL_PATH = config["embeddings"]["model_path"]
DEVICE = config["embeddings"].get("device", "cpu")
# VLLM_HOST = config["vllm"]["host"]
# VLLM_PORT = config["vllm"]["port"]
VLLM_BASE_URL=config["vllm"]["base_url"]
VLLM_API_KEY=config["vllm"]["api_key"]
VLLM_MODEL=config["vllm"]["model_name"]
VLLM_TEMPERATURE=config["vllm"]["temperature"]
VLLM_MAX_TOKENS=config["vllm"]["max_tokens"]
# print(FAISS_PATH,"=========================================================================")

# 实列化一个向量模型对象，这里使用的是下载到本地的向量模型，具体模型可以在config.yaml中配置
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL_PATH, model_kwargs={"device":"cuda:1"})
# 全局变量
vectorstore = None
retriever = None
# 定义文件索引之间的映射关系
file_index_map = {}

def load_or_create_vectorstore():
    global vectorstore, retriever, file_index_map

    if os.path.exists(FAISS_PATH):
        vectorstore = FAISS.load_local(FAISS_PATH, embeddings,allow_dangerous_deserialization=True)
        if not hasattr(vectorstore, "docstore") or vectorstore.docstore is None:
            vectorstore.docstore = InMemoryDocstore({})
            vectorstore.index_to_docstore_id = {}
        print("FAISS向量库已加载")
    else:
        # 初始化空向量库
        # dim = embeddings.embed_query("init").shape[0]
        # 初始化一个文本到向量模型，获得返回值
        dim = embeddings.embed_documents(["init"])[0]
        # 获取返回值的长度，即维度
        lendim=len(dim)
        vectorstore = FAISS(
            index=faiss.IndexFlatL2(lendim), # 定义向量数据库维度
            embedding_function=embeddings, #定义embeddings模型
            docstore=InMemoryDocstore({}),  #定义内存字典存储，通过id拿回原文   #docstore.add({"1": "这是第1个文档", "2": "这是第2个文档"})
            index_to_docstore_id={}  #index_to_docstore_id 是用来 做索引位置（FAISS 内部顺序） 和 文档 ID 的映射。 {0: "id1",1: "id2",2: "id3",...}
        )
        os.makedirs(FAISS_PATH, exist_ok=True)
        # 保存向量到本地
        vectorstore.save_local(FAISS_PATH)
        print("FAISS向量库已创建")
    
    # 初始化 retriever
    retriever = vectorstore.as_retriever(search_kwargs={"k":3})

    # 初始化 file_index_map
    file_index_map = {doc.metadata.get("source_file"): doc.metadata.get("file_hash") 
                      for doc in vectorstore.docstore._dict.values() if doc.metadata.get("source_file")}
    return vectorstore
#初始化向量数据实列
load_or_create_vectorstore()
# 文件哈希计算
def file_hash(file_path):
    hash_md5 = hashlib.md5()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()

# ----------------------------
# 请求体
class QueryRequest(BaseModel):
    query: str

# 上传 md 文件并增量更新向量库
@app.post("/upload_md")
async def upload_md(file: UploadFile = File(...)):
    global vectorstore, retriever, file_index_map
    temp_path = f"./temp_{file.filename}"
    with open(temp_path, "wb") as f:
        shutil.copyfileobj(file.file, f)
    # 对文件进行hash计算
    hash_val = file_hash(temp_path)
    # 拼接文件索引map路径
    file_path = os.path.join(FAISS_PATH,"file_index_map.pkl")
    # 如果本地已存有文件索引映射文件，那么就直接从本地文件加载数据，pickle.load
    if os.path.exists(file_path):
        with open(file_path, "rb") as f:
            file_index_map = pickle.load(f)
    # 初始化集合,判断文件是否在映射表中
    if file.filename not in file_index_map:
        # 没有就设置为空元组
        file_index_map[file.filename] = set()
    # 如果文件未修改，跳过
    if hash_val in file_index_map[file.filename]:
        os.remove(temp_path)
        return {"message": f"{file.filename} 未修改，跳过更新"}
    # 加载文档
    loader = TextLoader(temp_path, encoding="utf-8")
    docs = loader.load()
    # 文本切分
    text_splitter = RecursiveCharacterTextSplitter(
        separators =["\n\n","\n","."," ",""],
        #separators =["\n\n"],
        chunk_size = 512, #切片大小512个字符，建议值
        chunk_overlap = 128, #重叠部分128个字符，chunk_size的10%，20%
        length_function=len,)
    # split_docs 为列表
    split_docs = text_splitter.split_documents(docs)
    # doc_list=[]
    # 给每个切分好的块文件添加源数据metadata
    for doc in split_docs:
        # 源数据包含源文件和源文件对应的文件hash值
        doc.metadata["source_file"] = file.filename
        doc.metadata["file_hash"] = hash_val
    # 判断向量数据库中是否有对应文件名称向量数据，即对应文件的索引ID
    old_doc_ids = [doc_id for doc_id, doc in vectorstore.docstore._dict.items()
                   if doc.metadata.get("source_file") == file.filename]
    # 如果有，就删除向量数据库中对应的索引ID
    if old_doc_ids:
        vectorstore.delete(ids=old_doc_ids)

    # 增量加入向量库  对向量数据库中添加新的文件切片向量
    vectorstore.add_documents(split_docs)
    # 保存到本地
    vectorstore.save_local(FAISS_PATH)
    # 将新的文件和文件hash添加的文件映射表中
    file_index_map[file.filename].add(hash_val)
    #保存 page_info 或 file_index_map
    with open(os.path.join(FAISS_PATH, "file_index_map.pkl"), "wb") as f:
        pickle.dump(file_index_map, f)
    # 更新 retriever
    global retriever
    # 实列化向量数据库，匹配相关性最高的3个文档切片
    retriever = vectorstore.as_retriever(search_kwargs={"k":3})
    os.remove(temp_path)
    return {"message": f"{file.filename} 已上传/更新向量库", "added_chunks": len(split_docs)}


@app.post("/rag_query")
async def rag_query(req: QueryRequest):
    query = req.query
    docs = retriever.get_relevant_documents(query)
    context = "\n".join([doc.page_content for doc in docs])[:3000]
    prompt = f"""
        你是一个智能助手。
        任务：判断以下检索到的内容是否与用户的问题相关。
        规则：
            1. 只输出 '和问题相关' 或 '和问题不相关'。
            2. 如果内容包含具体操作步骤、说明、命令或者配置方式，即使包含标题或编号，也算“和问题相关”。
            3. 不允许输出其他任何文字、解释或总结。
        用户问题: {query}
        检索到的内容: {context}
        请严格回答：
    """
    relevance_answer = await query_vllm(prompt, VLLM_BASE_URL, VLLM_API_KEY, VLLM_MODEL,
                                        temperature=1e-5, max_tokens=512, strict=True)
    print(relevance_answer)
    
    # 3. 如果相关，则直接返回文档内容，而不是生成新答案
    if relevance_answer == "和问题相关":
        
        # 直接从知识库返回文本
        final_answer = "\n".join([doc.page_content for doc in docs])
        final_answer = "\n".join([line for line in final_answer.splitlines() if line.strip() and not set(line.strip()) <= {"🕹️"}])
        answer_prompt = f"""
            根据以下内容回答问题：
            {final_answer}

            问题：{query}

            要求：
            1. 严格只根据上述内容回答，请根据原文文件内容来学习回答，根据原文来，不要回答和原文不相干的内容，不要多余补充。
            2. 不要生成任何额外字符、表情或符号、不要自己扩充，回答要简洁明了
        """
        sources = [doc.metadata.get("source_file") for doc in docs]
        final_answer = await query_vllm(answer_prompt, VLLM_BASE_URL, VLLM_API_KEY,
                                       VLLM_MODEL, temperature=0.0, max_tokens=512)
        return {
            "relevance": "和问题相关",
            "answer": final_answer,
            "sources": sources
        }
    else:
        return {
            "relevance": "和问题不相关",
            "answer": "知识库中未找到相关内容",
            "sources": []
        }
########################################qwen-agent-rag能力调用###########################################################      
from qwen_agent.agents import Assistant  
import json5
import json,re
from qwen_agent.tools.base import BaseTool, register_tool
from typing import Dict, List
from qwen_agent.utils.output_beautify import typewriter_print

@register_tool('doc_search')
class DocSearch(BaseTool):
    description = "基于本地向量索引的文档检索工具，用于查询知识库。"
    parameters = [{
        "name": "query",
        "type": "string",
        "description": "用户的问题或需要检索的文本",
        "required": True
    }]

    def call(self, params: str, **kwargs) -> str:
        # 解析 LLM 传来的参数
        query = json5.loads(params)["query"]
        # 从索引检索
        docs = retriever.get_relevant_documents(query)
        result="\n".join([doc.page_content for doc in docs])
        print("=============20250903====================",result)
        return result
llm_cfg = {
    # Use the model service provided by DashScope:
    #'model': 'qwen-max-latest',
    #'model_type': 'qwen_dashscope',
    #'api_key': 'sk-c79dbd3f68a4434c95bd81cbf1b82bcf',
    # It will use the `DASHSCOPE_API_KEY' environment variable if 'api_key' is not set here.

    # Use a model service compatible with the OpenAI API, such as vLLM or Ollama:
    # 'model': 'Qwen2.5-7B-Instruct',
    ##################################################################################
    # "model": "/data/data_ldf/models/models--Qwen--Qwen2.5-3B/snapshots/3aab1f1954e9cc14eb9509a215f9e5ca08227a9b",   # 必须是你本地加载成功的模型名字
    "model": "/data/data_ldf/models/Qwen/models--Qwen--Qwen3-4B-Instruct-2507/snapshots/eb25fbe4f35f7147763bc24445679d1c00588d89",
    "model_type": "oai",
    'model_server': 'http://192.168.10.3:8000/v1',  # base_url, also known as api_base
    'api_key': '1233456',

    # (Optional) LLM hyperparameters for generation:
    'generate_cfg': {
        'top_p': 1
    }
}
system_instruction = """你是一个知识库问答助手。
当用户提问时，必须调用 doc_search 工具检索本地文档内容，并使用检索结果生成答案。
不要凭空回答，也不要跳过工具。
严格只根据检索内容回答，请根据原文文件内容来学习回答，根据原文来，不要回答和原文不相干的内容，不要多余补充。
不要生成任何额外字符、表情或符号、不要自己扩充，回答要简洁明了
"""
tools = [DocSearch()]
bot = Assistant(
    llm=llm_cfg,
    system_message=system_instruction,
    function_list=tools,
)
# 聊天历史（如果需要多轮对话，可以存储在内存或 Redis）
# 用于存储多轮会话（简单内存实现）
chat_sessions: Dict[str, List[Dict[str, str]]] = {}
class QueryRequestAgent(BaseModel):
    session_id: str   # 会话ID（前端传，标识用户）
    query: str
@app.post("/rag_agent_query")
async def rag_query(req: QueryRequestAgent):
    session_id = req.session_id
    query = req.query
#    # 如果是新会话，初始化对话历史
#     if session_id not in chat_sessions:
#         chat_sessions[session_id] = [
#             {"role": "system", "content": system_instruction}
#         ]
#     # 确保消息列表以用户消息开始（系统消息除外）
#     user_msg = {"role": "user", "content": query}
#     chat_sessions[session_id].append(user_msg)

    messages = [{'role': 'user', 'content': query}]
    response_plain_text = ''
    for response in bot.run(messages=messages):
        response_plain_text = typewriter_print(response, response_plain_text)
    match = re.search(r"\[ANSWER\](.*)", response_plain_text, re.S)
    if match:
        answer_text = match.group(1).strip()
        return {
        
            "relevance": "和问题相关",
            "answer": answer_text
        }
    else:
        answer_text = response_plain_text 
        return {
        
            "relevance": "和问题相关",
            "answer": answer_text
        }
    
